# zone_voronoi_candidates.py
# ============================================================
# 城市功能区识别 —— 候选点为真实 POI，使用 全局 Voronoi（不重叠）
# 输出：output_zones.shp, output_centers.shp, zone_features.csv, zone_map.html
# ============================================================

import geopandas as gpd
import pandas as pd
import numpy as np
import math
import warnings
import os
from shapely.geometry import Point, MultiPoint, Polygon, box
from shapely.ops import unary_union
from sklearn.cluster import KMeans,DBSCAN
import folium
from folium.plugins import MarkerCluster
from tqdm import tqdm

# optional imports
USE_SHAPELY_VORONOI = False
USE_SCIPY_VORONOI = False
try:
    # shapely >= 2.0
    from shapely.ops import voronoi_diagram
    USE_SHAPELY_VORONOI = True
except Exception:
    try:
        from scipy.spatial import Voronoi
        USE_SCIPY_VORONOI = True
    except Exception:
        pass

# encoding detector
try:
    import chardet
except Exception:
    chardet = None

warnings.filterwarnings("ignore", category=UserWarning)
warnings.filterwarnings("ignore", category=FutureWarning)


# ---------------------------
# Helpers
# ---------------------------
def detect_encoding(path, default="utf-8"):
    if chardet is None:
        return default
    with open(path, "rb") as f:
        raw = f.read(100000)
    res = chardet.detect(raw)
    enc = res.get("encoding") or default
    return enc


def safe_read_csv(path):
    """
    Read CSV with encoding detection, fallback strategies.
    Returns pandas.DataFrame
    """
    # try direct utf-8 first (fast path)
    try:
        return pd.read_csv(path, encoding="utf-8")
    except Exception:
        pass

    # detect encoding
    enc = detect_encoding(path)
    try:
        return pd.read_csv(path, encoding=enc)
    except Exception:
        # last resort: open with errors='ignore' and let pandas parse file object (older pandas accept file-like)
        try:
            with open(path, "r", encoding=enc, errors="ignore") as f:
                return pd.read_csv(f)
        except Exception:
            # try common encodings
            for e in ["gbk", "latin1", "utf-8-sig"]:
                try:
                    return pd.read_csv(path, encoding=e)
                except Exception:
                    continue
    raise RuntimeError(f"无法读取 CSV 文件: {path}")


def get_union(geoms):
    """
    Prefer union_all if available (faster for GeoDataFrame), otherwise unary_union.
    Accepts GeoSeries or list of geometries.
    """
    try:
        # GeoSeries has union_all
        return geoms.union_all()
    except Exception:
        # fallback
        if hasattr(geoms, "unary_union"):
            return geoms.unary_union
        else:
            return unary_union(list(geoms))


# ---------------------------
# Voronoi helpers
# ---------------------------
def make_finite_voronoi_polygons(points_xy, bbox=None):
    """
    Create finite voronoi polygons from Nx2 numpy array of points.
    Use shapely.voronoi_diagram if available, otherwise SciPy Voronoi fallback.
    Returns list of shapely polygons (not ordered).
    """
    if len(points_xy) == 0:
        return []

    if USE_SHAPELY_VORONOI:
        mpts = MultiPoint([Point(xy) for xy in points_xy])
        # envelope argument will bound the diagram
        env = None
        if bbox is not None:
            env = box(*bbox)
        vor = voronoi_diagram(mpts, envelope=env)
        polys = [g for g in vor.geoms if isinstance(g, Polygon)]
        return polys

    # SciPy fallback
    from scipy.spatial import Voronoi
    vor = Voronoi(points_xy)

    # bounding box
    if bbox is None:
        pt_min = points_xy.min(axis=0)
        pt_max = points_xy.max(axis=0)
        span = pt_max - pt_min
        expand = span.max() * 10.0 + 1000
        minx, miny = pt_min[0] - expand, pt_min[1] - expand
        maxx, maxy = pt_max[0] + expand, pt_max[1] + expand
    else:
        minx, miny, maxx, maxy = bbox

    bbox_poly = Polygon([(minx, miny), (minx, maxy), (maxx, maxy), (maxx, miny)])

    regions = []
    for point_idx, region_idx in enumerate(vor.point_region):
        region = vor.regions[region_idx]
        if not region or -1 in region:
            # build polygon from ridges for infinite regions (approximate)
            # collect ridge vertices for this point
            verts = []
            for ridge_point_indices, ridge_vertices in zip(vor.ridge_points, vor.ridge_vertices):
                if point_idx in ridge_point_indices:
                    # ridge_vertices may contain -1
                    for v in ridge_vertices:
                        if v != -1:
                            verts.append(tuple(vor.vertices[v]))
            if len(verts) >= 3:
                poly = Polygon(verts).convex_hull
            else:
                poly = bbox_poly
        else:
            verts = [tuple(vor.vertices[i]) for i in region]
            poly = Polygon(verts)
        poly = poly.intersection(bbox_poly)
        regions.append(poly)
    return regions


def build_global_voronoi(points_gdf_proj, clip_poly=None):
    """
    Build Voronoi polygons from points_gdf_proj (projected CRS).
    Returns GeoDataFrame with geometry list in same CRS, order not guaranteed.
    """
    pts = np.array([[p.x, p.y] for p in points_gdf_proj.geometry])
    if len(pts) == 0:
        return gpd.GeoDataFrame(columns=["geometry"], crs=points_gdf_proj.crs)

    # bbox
    minx, miny, maxx, maxy = pts[:, 0].min(), pts[:, 1].min(), pts[:, 0].max(), pts[:, 1].max()
    pad = max(maxx - minx, maxy - miny) * 5 + 1000
    bbox = (minx - pad, miny - pad, maxx + pad, maxy + pad)

    polys = make_finite_voronoi_polygons(pts, bbox=bbox)

    # clip to clip_poly if provided
    if clip_poly is not None:
        polys = [poly.intersection(clip_poly) if (poly is not None and not poly.is_empty) else poly for poly in polys]

    # Filter valid non-empty polys
    polys = [poly for poly in polys if poly is not None and (not poly.is_empty) and poly.is_valid]

    # Return GeoDataFrame
    vor_gdf = gpd.GeoDataFrame(geometry=polys, crs=points_gdf_proj.crs)
    return vor_gdf


# ---------------------------
# Pipeline functions
# ---------------------------
def read_inputs(boundary_path, road_path, poi_path, luckin_path=None):
    print("\n--- 读取输入 ----")
    boundary_gdf = gpd.read_file(boundary_path)
    roads_gdf = gpd.read_file(road_path)

    # try reading poi as vector first
    poi_gdf = None
    try:
        poi_gdf = gpd.read_file(poi_path)
    except Exception:
        df = safe_read_csv(poi_path)
        # detect lon/lat columns
        lon_col = next((c for c in df.columns if "lon" in c.lower() or "经度" in c), None)
        lat_col = next((c for c in df.columns if "lat" in c.lower() or "纬度" in c), None)
        if lon_col and lat_col:
            poi_gdf = gpd.GeoDataFrame(df, geometry=gpd.points_from_xy(df[lon_col], df[lat_col]), crs="EPSG:4326")
        else:
            raise ValueError("POI 文件无法解析为 GeoDataFrame：缺少经纬度列")

    luckin_gdf = None
    if luckin_path:
        if luckin_path.lower().endswith(".csv"):
            df = safe_read_csv(luckin_path)
            lon_col = next((c for c in df.columns if "lon" in c.lower() or "经度" in c), None)
            lat_col = next((c for c in df.columns if "lat" in c.lower() or "纬度" in c), None)
            if lon_col and lat_col:
                luckin_gdf = gpd.GeoDataFrame(df, geometry=gpd.points_from_xy(df[lon_col], df[lat_col]), crs="EPSG:4326")
        else:
            try:
                luckin_gdf = gpd.read_file(luckin_path)
            except Exception:
                luckin_gdf = None

    # ensure CRS
    for gdf in [boundary_gdf, roads_gdf, poi_gdf, luckin_gdf]:
        if gdf is not None and getattr(gdf, "crs", None) is None:
            print("⚠️ 输入文件缺少 CRS，假定 EPSG:4326")
            gdf.set_crs(epsg=4326, inplace=True)

    print(f"边界: {len(boundary_gdf)} 要素, 道路: {len(roads_gdf)} 条, POI: {len(poi_gdf)} 条, 瑞幸门店: {len(luckin_gdf) if luckin_gdf is not None else 0}")
    return boundary_gdf, roads_gdf, poi_gdf, luckin_gdf


def select_candidate_points(poi_gdf, k_per_type=5):
    """
    For each category in poi_gdf, do KMeans and select nearest real POI to each centroid.
    Return GeoDataFrame of candidate points (projected EPSG:3857).
    """
    if "category" not in poi_gdf.columns:
        print("⚠️ POI 缺失 'category' 列，全部作为 'Unknown'")
        poi_gdf = poi_gdf.copy()
        poi_gdf["category"] = "Unknown"

    poi_proj = poi_gdf.to_crs(epsg=3857)
    centers = []

    for cat, group in poi_proj.groupby("category"):
        group = group.copy()
        n = len(group)
        if n == 0:
            continue
        k = min(k_per_type, n)
        if k <= 0:
            continue

        if n == 1:
            centers.append({"category": cat, "geometry": group.geometry.iloc[0]})
            continue

        coords = np.vstack((group.geometry.x.values, group.geometry.y.values)).T
        # handle cases where k==1
        if k == 1:
            centroid = coords.mean(axis=0)
            dists = np.hypot(coords[:, 0] - centroid[0], coords[:, 1] - centroid[1])
            idx = int(np.argmin(dists))
            row = group.iloc[idx]
            centers.append({"category": cat, "geometry": row.geometry})
            continue

        kmeans = KMeans(n_clusters=k, n_init="auto", random_state=42)
        labels = kmeans.fit_predict(coords)
        centroids = kmeans.cluster_centers_

        for i in range(k):
            member_mask = (labels == i)
            members = group.iloc[np.where(member_mask)[0]]
            if len(members) == 0:
                continue
            centroid = centroids[i]
            dists = np.hypot(members.geometry.x - centroid[0], members.geometry.y - centroid[1])
            nearest_idx = dists.idxmin()
            nearest_row = members.loc[nearest_idx]
            centers.append({"category": cat, "geometry": nearest_row.geometry})

    centers_gdf = gpd.GeoDataFrame(centers, crs="EPSG:3857")
    # drop duplicates (same point selected multiple times)
    centers_gdf = centers_gdf.drop_duplicates(subset=["geometry"])
    centers_gdf = centers_gdf.reset_index(drop=True)
    return centers_gdf


def generate_global_voronoi_from_candidates(centers_gdf_proj, poi_proj, boundary_proj):
    """
    Build global Voronoi from all candidate points (projected), clip to boundary,
    assign each polygon to nearest candidate (inherit category), and optionally clip
    further to category convex hull to avoid far-reaching cells.
    """
    if len(centers_gdf_proj) == 0:
        return gpd.GeoDataFrame(columns=["category", "geometry"]), centers_gdf_proj

    # build voronoi polygons
    # clip_poly base: boundary union
    clip_base = get_union(boundary_proj.geometry)

    vor_gdf = build_global_voronoi(centers_gdf_proj, clip_poly=clip_base)

    # assign each voronoi polygon to nearest candidate point (by centroid nearest)
    assigned_records = []
    centers_pts = np.array([[p.x, p.y] for p in centers_gdf_proj.geometry])
    centers_cat = centers_gdf_proj["category"].values

    for poly in vor_gdf.geometry:
        if poly.is_empty:
            continue
        rep = poly.representative_point()
        dists = np.linalg.norm(centers_pts - np.array([rep.x, rep.y]), axis=1)
        i = int(np.argmin(dists))
        cat = centers_cat[i]
        assigned_records.append({"category": cat, "geometry": poly})

    assigned_gdf = gpd.GeoDataFrame(assigned_records, crs=centers_gdf_proj.crs)

    # Optionally clip each assigned polygon to category convex hull ∩ boundary to reduce overreach
    final_polys = []
    for _, row in assigned_gdf.iterrows():
        cat = row["category"]
        poly = row.geometry
        # build cat hull
        group_pois = poi_proj[poi_proj["category"] == cat]
        if len(group_pois) > 0:
            try:
                hull = get_union(group_pois.geometry).convex_hull.buffer(500)
                poly_clipped = poly.intersection(hull).intersection(clip_base)
            except Exception:
                poly_clipped = poly.intersection(clip_base)
        else:
            poly_clipped = poly.intersection(clip_base)
        # keep valid poly
        if poly_clipped.is_empty or (not poly_clipped.is_valid):
            poly_clipped = poly.intersection(clip_base)
        final_polys.append({"category": cat, "geometry": poly_clipped})

    final_gdf = gpd.GeoDataFrame(final_polys, crs=centers_gdf_proj.crs)
    # remove empty
    final_gdf = final_gdf[~final_gdf.geometry.is_empty & final_gdf.geometry.is_valid].reset_index(drop=True)
    return final_gdf, centers_gdf_proj


def compute_zone_stats(zones_gdf_proj, poi_gdf_proj, luckin_gdf_proj=None):
    records = []
    for idx, row in zones_gdf_proj.reset_index(drop=True).iterrows():
        poly = row.geometry
        poi_in = poi_gdf_proj[poi_gdf_proj.geometry.within(poly)]
        area_km2 = poly.area / 1e6
        rec = {
            "zone_id": idx,
            "category": row.get("category", ""),
            "poi_count": len(poi_in),
            "area_km2": round(area_km2, 4),
            "poi_density_per_km2": round(len(poi_in) / area_km2 if area_km2 > 0 else 0, 3)
        }
        if luckin_gdf_proj is not None:
            rec["luckin_count"] = int(len(luckin_gdf_proj[luckin_gdf_proj.geometry.within(poly)]))
        records.append(rec)
    return pd.DataFrame(records)


def add_legend(fmap, color_map, title="Category"):
    """
    Adds a small HTML legend to folium map.
    color_map: dict {category: color}
    """
    legend_items = ""
    for cat, col in color_map.items():
        legend_items += f"""
            <li style="list-style:none;margin-bottom:4px;">
              <span style="display:inline-block;width:12px;height:12px;background:{col};margin-right:6px;border:1px solid #999;"></span>{cat}
            </li>
        """
    html = f"""
    <div style="
        position: fixed;
        bottom: 30px;
        left: 10px;
        z-index:9999;
        background:white;
        padding:10px;
        border:1px solid #ccc;
        box-shadow:2px 2px 6px rgba(0,0,0,0.2);
        font-size:12px;
    ">
      <strong>{title}</strong>
      <ul style="padding-left:8px;margin:6px 0 0 0;">
        {legend_items}
      </ul>
    </div>
    """
    from folium import IFrame
    fmap.get_root().html.add_child(folium.Element(html))


def visualize_to_html(boundary_gdf, zones_gdf_proj, centers_gdf_proj, poi_gdf, luckin_gdf=None, out_html="zone_map.html"):
    # convert to WGS84
    boundary_4326 = boundary_gdf.to_crs(epsg=4326)
    zones_4326 = zones_gdf_proj.to_crs(epsg=4326)
    centers_4326 = centers_gdf_proj.to_crs(epsg=4326)
    poi_4326 = poi_gdf.to_crs(epsg=4326)
    luckin_4326 = luckin_gdf.to_crs(epsg=4326) if (luckin_gdf is not None) else None

    # center map
    if len(poi_4326) > 0:
        center_lat = float(poi_4326.geometry.y.mean())
        center_lon = float(poi_4326.geometry.x.mean())
    else:
        center = boundary_4326.unary_union.centroid
        center_lat, center_lon = center.y, center.x

    m = folium.Map(location=[center_lat, center_lon], zoom_start=13, tiles="cartodbpositron")

    # boundary
    folium.GeoJson(boundary_4326, name="Boundary", style_function=lambda x: {"color": "black", "weight": 1, "fill": False}).add_to(m)

    # color map for categories
    cats = list(zones_gdf_proj["category"].unique())
    # choose colors
    import matplotlib.pyplot as plt
    cmap = plt.get_cmap("tab20")
    colors = [mcolors_to_hex(cmap(i / max(1, len(cats)-1))) for i in range(len(cats))]
    color_map = {cat: colors[i % len(colors)] for i, cat in enumerate(cats)}

    # add zones
    for _, row in zones_4326.iterrows():
        cat = row["category"]
        col = color_map.get(cat, "#3186cc")
        folium.GeoJson(
            row.geometry,
            style_function=lambda feat, color=col: {"fillColor": color, "color": color, "weight": 1, "fillOpacity": 0.35},
            tooltip=f"Category: {cat}"
        ).add_to(m)

    # add centers
    for _, row in centers_4326.iterrows():
        folium.CircleMarker([row.geometry.y, row.geometry.x], radius=5, color="red", fill=True, fill_opacity=0.9,
                            popup=f"Category: {row.get('category','') if 'category' in row else ''}").add_to(m)

    # add poi cluster
    mc = MarkerCluster(name="POI").add_to(m)
    for _, row in poi_4326.iterrows():
        folium.CircleMarker([row.geometry.y, row.geometry.x], radius=3, color="black", fill=True, fill_opacity=0.7,
                            popup=str(row.get("category", ""))).add_to(mc)

    # add luckin points
    if luckin_4326 is not None:
        for _, row in luckin_4326.iterrows():
            folium.CircleMarker([row.geometry.y, row.geometry.x], radius=5, color="blue", fill=True, fill_opacity=0.9,
                                popup="Luckin").add_to(m)

    # add legend
    add_legend(m, color_map, title="Category")

    folium.LayerControl().add_to(m)
    m.save(out_html)
    print(f"✅ HTML 地图已生成: {out_html}")


def mcolors_to_hex(rgb_tuple):
    """Convert matplotlib rgba to hex."""
    try:
        # rgb_tuple could be (r,g,b,a) or (r,g,b)
        r, g, b = rgb_tuple[:3]
        return '#%02x%02x%02x' % (int(r * 255), int(g * 255), int(b * 255))
    except Exception:
        return "#3186cc"


# ---------------------------
# Main
# ---------------------------
def main():
    # paths (modify if needed)
    boundary_path = "仓山区最终完整边界.shp"
    road_path = "仓山区道路信息.shp"
    poi_path = "poi仓山区.shp"  # shp or csv (with lon/lat)
    luckin_path = "福州市仓山区瑞幸咖啡坐标.csv"  # optional csv
    k_per_type = 5

    # 1. read
    boundary_gdf, roads_gdf, poi_gdf, luckin_gdf = read_inputs(boundary_path, road_path, poi_path, luckin_path)

    # 2. select candidate points (projected)
    centers_gdf_proj = select_candidate_points(poi_gdf, k_per_type=k_per_type)

    # 3. build global voronoi (projected) and assign categories
    poi_proj = poi_gdf.to_crs(epsg=3857)
    boundary_proj = boundary_gdf.to_crs(epsg=3857)
    zones_gdf_proj, centers_gdf_proj = generate_global_voronoi_from_candidates(centers_gdf_proj, poi_proj, boundary_proj)

    # 4. compute stats
    luckin_proj = luckin_gdf.to_crs(epsg=3857) if (luckin_gdf is not None) else None
    stats_df = compute_zone_stats(zones_gdf_proj, poi_proj, luckin_proj)

    # 5. export
    try:
        zones_out = zones_gdf_proj.to_crs(epsg=4326)
        centers_out = centers_gdf_proj.to_crs(epsg=4326)
    except Exception:
        zones_out = zones_gdf_proj.copy()
        centers_out = centers_gdf_proj.copy()

    zones_out.to_file("output_zones.shp", encoding="utf-8")
    centers_out.to_file("output_centers.shp", encoding="utf-8")
    stats_df.to_csv("zone_features.csv", index=False, encoding="utf-8-sig")
    print("✅ 输出 shapefile 与 CSV 完成：output_zones.shp, output_centers.shp, zone_features.csv")

    # 6. visualize
    visualize_to_html(boundary_gdf, zones_gdf_proj, centers_gdf_proj, poi_gdf, luckin_gdf, out_html="zone_map.html")
    print("\n--- 全流程完成 ---")


if __name__ == "__main__":
    main()
